/*
Copyright (C) 2022-2024 Inspur Corp. All Rights Reserved.

SPDX-License-Identifier: Apache-2.0
*/
package gm

import (
	"crypto/rand"
	"did-service/common/file"
	"encoding/binary"
	"encoding/hex"
	"io"
	"os"

	"github.com/pkg/errors"
	"github.com/tjfoc/gmsm/sm4"
)

const (
	FrameSize     = 2048
	DefaultSm4Key = "5598662321a0af1eafa6e2591858aafd"
	SM4Name       = "SM4"
)

// 生成SM4秘钥，秘钥长度16字节，输出为16进制编码字符串
func GenerateSm4Key() string {
	bytes, _ := getRandomBytes(16)
	result := hex.EncodeToString(bytes)
	return result
}

// 加密私钥
func EncryptPrvKey(prvKeyHex string) string {
	bytes, err := hex.DecodeString(prvKeyHex)
	if err != nil {
		log.Errorf("encrypt private key, decode key error: %v", err)
		return ""
	}
	encryptedBytes := Sm4Encrypt(DefaultSm4Key, bytes)
	result := hex.EncodeToString(encryptedBytes)
	return result
}

// 解密私钥
func DecryptPrvKey(decrytedPrvKeyHex string) string {
	bytes, err := hex.DecodeString(decrytedPrvKeyHex)
	if err != nil {
		log.Errorf("decrypt private key, decode key error: %v", err)
		return ""
	}
	decryptedBytes := Sm4Decrypt(DefaultSm4Key, bytes)
	result := hex.EncodeToString(decryptedBytes)
	return result
}

// SM4加密
func Sm4Encrypt(keyHex string, data []byte) []byte {
	key, err := hex.DecodeString(keyHex)
	if err != nil {
		log.Errorf("sm4 encrypt, decode key error: %v", err)
		return nil
	}
	bytes, err := sm4.Sm4OFB(key, data, true)
	if err != nil {
		log.Errorf("sm4 encrypt error: %v", err)
		return nil
	}
	return bytes
}

// SM4解密
func Sm4Decrypt(keyHex string, encryptedData []byte) []byte {
	key, err := hex.DecodeString(keyHex)
	if err != nil {
		log.Errorf("sm4 decrypt, decode key error: %v", err)
		return nil
	}
	bytes, err := sm4.Sm4OFB(key, encryptedData, false)
	if err != nil {
		log.Errorf("sm4 decrypt error: %v", err)
		return nil
	}
	return bytes
}

// SM4加密文件
func Sm4EncryptFile(keyHex string, srcPath string, destPath string) error {
	key, err := hex.DecodeString(keyHex)
	if err != nil {
		return err
	}

	srcFile, err := os.Open(srcPath)
	if err != nil {
		return err
	}
	defer srcFile.Close()

	exists := file.FileExists(destPath)
	if exists {
		return errors.New("dest file exist!")
	}
	destFile, err := os.Create(destPath)
	if err != nil {
		return err
	}
	defer destFile.Close()

	frameBuf := make([]byte, FrameSize)
	sizeBuf := make([]byte, 2)
	for {
		n, err := srcFile.Read(frameBuf)
		if err != nil {
			if err == io.EOF {
				break
			}
			return err
		}
		if n <= 0 {
			break
		}
		bytes, err := sm4.Sm4OFB(key, frameBuf[:n], true)
		binary.BigEndian.PutUint16(sizeBuf[:], uint16(len(bytes)))
		_, err = destFile.Write(sizeBuf)
		if err != nil {
			return err
		}
		_, err = destFile.Write(bytes)
		if err != nil {
			return err
		}
	}
	return nil
}

// SM4解密文件
func Sm4DecryptFile(keyHex string, srcPath string, destPath string) error {
	key, err := hex.DecodeString(keyHex)
	if err != nil {
		return err
	}

	srcFile, err := os.Open(srcPath)
	if err != nil {
		return err
	}
	defer srcFile.Close()

	exists := file.FileExists(destPath)
	if exists {
		return errors.New("dest file exist!")
	}
	destFile, err := os.Create(destPath)
	if err != nil {
		return err
	}
	defer destFile.Close()

	sizeBuf := make([]byte, 2)

	for {
		n, err := srcFile.Read(sizeBuf)
		if err != nil {
			if err == io.EOF {
				break
			}
			return err
		}
		if n <= 0 {
			break
		}
		encryptedFrameSize := binary.BigEndian.Uint16(sizeBuf)

		buf := make([]byte, encryptedFrameSize)
		n, err = srcFile.Read(buf)
		if err != nil {
			if err == io.EOF {
				break
			}
			return err
		}
		if n <= 0 {
			break
		}
		bytes, err := sm4.Sm4OFB(key, buf[:n], false)
		destFile.Write(bytes)
	}

	return nil
}

func getRandomBytes(len int) ([]byte, error) {
	if len < 0 {
		return nil, errors.New("Len must be larger than 0")
	}
	buffer := make([]byte, len)
	n, err := rand.Read(buffer)
	if err != nil {
		return nil, err
	}
	if n != len {
		return nil, errors.Errorf("Buffer not filled. Requested [%d], got [%d]", len, n)
	}
	return buffer, nil
}
